#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自动交易程序 - 数据库修复增强版

此版本修复了数据库表结构不匹配导致的交易记录无法插入问题，
优化了数据库操作逻辑和错误处理机制。
"""

import time
import os
import sys
import logging
import pymysql
import sqlite3
from datetime import datetime, timedelta
import random
import math
from tqsdk import TqApi, TqKq, TqAuth
from tqsdk.ta import MA
from rich.console import Console
from rich.theme import Theme

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("trading_log.log", encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("AutoTrader")

# 定义控制台输出，更明亮显眼颜色主题
custom_theme = Theme({
    "info": "bright_cyan",
    "warning": "bright_yellow",
    "error": "bright_red",
    "success": "bright_green",
    "trade": "bold bright_magenta"
})
console = Console(theme=custom_theme)

# MySQL数据库配置
MYSQL_CONFIG = {
    'host': '172.28.196.123',
    'user': 'remote_admin',
    'password': 'alibaba999',
    'database': 'trade_monitor',
    'port': 3306
}

# SQLite数据库配置
SQLITE_DB_PATH = os.path.join(os.path.dirname(__file__), 'trade_monitor.db')

# 全局配置
POSITION_FUND_RATIO = 0.2  # 仓位资金比例
MAX_RETRY_COUNT = 3  # 数据库操作最大重试次数
RETRY_INTERVAL = 2  # 重试间隔(秒)

# 监控的品种列表
MONITOR_SYMBOLS = [
    "SHFE.rb2501", "SHFE.ru2501", "SHFE.al2501", "SHFE.cu2501",
    "DCE.i2501", "DCE.m2501", "DCE.p2501", "DCE.y2501",
    "CZCE.MA2501", "CZCE.SR2501", "CZCE.CF2501", "CZCE.RM2501"
]


def get_config_value(key, default_value=None):
    """从配置表获取配置值"""
    try:
        # 优先从MySQL获取
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        cursor.execute("SELECT value FROM config WHERE `key` = %s", (key,))
        result = cursor.fetchone()
        cursor.close()
        conn.close()
        if result:
            return result[0]
    except Exception as e:
        logger.error(f"从MySQL获取配置出错: {str(e)}")
        
    # MySQL失败时从SQLite获取
    try:
        conn = sqlite3.connect(SQLITE_DB_PATH)
        cursor = conn.cursor()
        cursor.execute("SELECT value FROM config WHERE `key` = ?", (key,))
        result = cursor.fetchone()
        cursor.close()
        conn.close()
        if result:
            return result[0]
    except Exception as e:
        logger.error(f"从SQLite获取配置出错: {str(e)}")
    
    return default_value


def connect_to_tqsdk():
    """连接到TQSDK，支持重试机制"""
    max_retries = 3
    retry_count = 0
    
    while retry_count < max_retries:
        try:
            # 使用模拟账户进行连接测试
            api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba"), debug=False)
            console.print("成功连接到TQSDK！", style="success")
            logger.info("成功连接到TQSDK")
            return api
        except Exception as e:
            retry_count += 1
            console.print(f"连接TQSDK失败 (尝试 {retry_count}/{max_retries}): {str(e)}", style="error")
            logger.error(f"连接TQSDK失败 (尝试 {retry_count}/{max_retries}): {str(e)}")
            
            if retry_count < max_retries:
                console.print(f"{RETRY_INTERVAL}秒后重试...", style="warning")
                time.sleep(RETRY_INTERVAL)
            else:
                console.print("达到最大重试次数，使用模拟模式继续运行", style="warning")
                logger.warning("达到最大重试次数，使用模拟模式继续运行")
                # 返回None表示需要使用模拟模式
                return None


def connect_to_mysql():
    """连接到MySQL数据库，支持重试机制"""
    retry_count = 0
    
    while retry_count < MAX_RETRY_COUNT:
        try:
            conn = pymysql.connect(**MYSQL_CONFIG)
            logger.info("成功连接到MySQL数据库")
            return conn
        except Exception as e:
            retry_count += 1
            logger.error(f"连接MySQL数据库失败 (尝试 {retry_count}/{MAX_RETRY_COUNT}): {str(e)}")
            
            if retry_count < MAX_RETRY_COUNT:
                time.sleep(RETRY_INTERVAL)
            else:
                logger.warning("达到最大重试次数，无法连接到MySQL数据库")
                return None


def connect_to_sqlite():
    """连接到SQLite数据库"""
    try:
        # 确保SQLite数据库目录存在
        os.makedirs(os.path.dirname(SQLITE_DB_PATH), exist_ok=True)
        
        conn = sqlite3.connect(SQLITE_DB_PATH)
        logger.info("成功连接到SQLite数据库")
        return conn
    except Exception as e:
        logger.error(f"连接SQLite数据库失败: {str(e)}")
        return None


def ensure_db_tables_exists():
    """确保数据库表存在"""
    # 确保SQLite表存在
    try:
        sqlite_conn = connect_to_sqlite()
        if sqlite_conn:
            cursor = sqlite_conn.cursor()
            
            # 创建trade_records表
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS trade_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                product_name TEXT NOT NULL,
                direction TEXT NOT NULL,
                offset TEXT NOT NULL,
                volume INTEGER NOT NULL,
                price REAL NOT NULL,
                reason TEXT NOT NULL,
                timestamp TEXT NOT NULL
            )
            """)
            
            # 创建config表
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS config (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                `key` TEXT NOT NULL UNIQUE,
                `value` TEXT NOT NULL
            )
            """)
            
            # 初始化默认配置
            default_configs = [
                ('open_auto', 'on'),
                ('math_open_lots', 'off'),
                ('position_fund_ratio', '0.2'),
                ('stop_loss_enabled', 'on'),
                ('monitor_interval', '30')
            ]
            
            for key, value in default_configs:
                cursor.execute("SELECT COUNT(*) FROM config WHERE `key` = ?", (key,))
                if cursor.fetchone()[0] == 0:
                    cursor.execute("INSERT INTO config (`key`, `value`) VALUES (?, ?)", (key, value))
            
            sqlite_conn.commit()
            cursor.close()
            sqlite_conn.close()
    except Exception as e:
        logger.error(f"初始化SQLite表结构时出错: {str(e)}")


def insert_trade_record(record_data):
    """插入交易记录到数据库，支持自动切换数据库"""
    # 准备插入数据
    record = {
        'product_name': record_data.get('product_name', '未知'),
        'direction': record_data.get('direction', 'UNKNOWN'),
        'offset': record_data.get('offset', 'OPEN'),
        'volume': record_data.get('volume', 0),
        'price': record_data.get('price', 0.0),
        'reason': record_data.get('reason', '未指定'),
        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }
    
    # 先尝试插入到MySQL
    mysql_success = False
    try:
        mysql_conn = connect_to_mysql()
        if mysql_conn:
            cursor = mysql_conn.cursor()
            
            # 准备SQL语句，注意包含所有必要字段
            sql = """
            INSERT INTO trade_records 
            (product_name, direction, offset, volume, price, reason, timestamp)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            """
            
            # 执行插入
            cursor.execute(sql, (
                record['product_name'], record['direction'], record['offset'],
                record['volume'], record['price'], record['reason'], record['timestamp']
            ))
            mysql_conn.commit()
            last_id = cursor.lastrowid
            
            cursor.close()
            mysql_conn.close()
            mysql_success = True
            logger.info(f"成功向MySQL插入交易记录，ID: {last_id}")
    except Exception as e:
        logger.error(f"向MySQL插入交易记录失败: {str(e)}")
    
    # MySQL失败时，插入到SQLite
    if not mysql_success:
        try:
            sqlite_conn = connect_to_sqlite()
            if sqlite_conn:
                cursor = sqlite_conn.cursor()
                
                # 准备SQL语句
                sql = """
                INSERT INTO trade_records 
                (product_name, direction, offset, volume, price, reason, timestamp)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """
                
                # 执行插入
                cursor.execute(sql, (
                    record['product_name'], record['direction'], record['offset'],
                    record['volume'], record['price'], record['reason'], record['timestamp']
                ))
                sqlite_conn.commit()
                last_id = cursor.lastrowid
                
                cursor.close()
                sqlite_conn.close()
                logger.info(f"成功向SQLite插入交易记录，ID: {last_id}")
                return True
        except Exception as e:
            logger.error(f"向SQLite插入交易记录失败: {str(e)}")
    
    return mysql_success


def check_position_and_orders(api, symbol):
    """检查指定品种的持仓和订单"""
    try:
        # 获取持仓
        positions = api.get_position(symbol)
        position_volume = positions.get('volume_long', 0) - positions.get('volume_short', 0)
        
        # 获取未完成订单
        orders = api.get_order(symbol)
        active_orders = [o for o in orders.values() if o['status'] != 'FINISHED']
        
        return {
            'position_volume': position_volume,
            'active_orders_count': len(active_orders),
            'active_orders': active_orders
        }
    except Exception as e:
        logger.error(f"检查{symbol}持仓和订单时出错: {str(e)}")
        return {
            'position_volume': 0,
            'active_orders_count': 0,
            'active_orders': []
        }


def stop_loss(api, symbol, position_info):
    """止损逻辑"""
    try:
        if get_config_value('stop_loss_enabled', 'on') != 'on':
            return False
        
        # 获取当前行情
        quote = api.get_quote(symbol)
        
        # 简单的止损逻辑示例
        # 这里应该根据实际策略实现更复杂的止损逻辑
        
        if position_info['position_volume'] > 0:
            # 多头止损
            # 可以实现跟踪止损、固定比例止损等
            pass
        elif position_info['position_volume'] < 0:
            # 空头止损
            # 可以实现跟踪止损、固定比例止损等
            pass
        
        return True
    except Exception as e:
        logger.error(f"执行{symbol}止损逻辑时出错: {str(e)}")
        return False


def calculate_lots(api, symbol, direction):
    """计算开仓手数"""
    try:
        # 获取当前行情
        quote = api.get_quote(symbol)
        last_price = quote.get('last_price', 0)
        
        if last_price <= 0:
            logger.warning(f"{symbol}价格无效: {last_price}")
            return 1  # 返回默认手数
        
        # 获取合约信息
        instrument = api.get_instrument(symbol)
        price_tick = instrument.get('price_tick', 1)
        volume_multiple = instrument.get('volume_multiple', 1)
        
        # 根据账户资金计算手数
        # 这里使用简单的计算方式，实际应该根据账户资金和风险偏好计算
        fund_ratio = float(get_config_value('position_fund_ratio', '0.2'))
        
        # 假设账户中有100万资金用于交易
        account_fund = 1000000
        
        # 计算保证金
        # 注意：这里需要根据实际的保证金比例计算
        margin_ratio = 0.1  # 假设保证金比例为10%
        margin_per_lot = last_price * volume_multiple * margin_ratio
        
        # 计算可开手数
        max_lots = int(account_fund * fund_ratio / margin_per_lot)
        
        # 确保至少开1手
        lots = max(1, min(10, max_lots))  # 限制最大手数为10
        
        return lots
    except Exception as e:
        logger.error(f"计算{symbol}开仓手数时出错: {str(e)}")
        return 1  # 返回默认手数


def open_position_non_shfe(api, symbol, direction):
    """开仓逻辑（非上海交易所品种）"""
    try:
        # 检查当前持仓和订单
        position_info = check_position_and_orders(api, symbol)
        
        # 如果有未完成订单，不重复开仓
        if position_info['active_orders_count'] > 0:
            logger.info(f"{symbol}有未完成订单，暂时不开仓")
            return False
        
        # 根据方向检查持仓
        if direction == 'BUY' and position_info['position_volume'] > 0:
            logger.info(f"{symbol}已有多头持仓，暂时不开新仓")
            return False
        if direction == 'SELL' and position_info['position_volume'] < 0:
            logger.info(f"{symbol}已有空头持仓，暂时不开新仓")
            return False
        
        # 获取当前行情
        quote = api.get_quote(symbol)
        last_price = quote.get('last_price', 0)
        
        if last_price <= 0:
            logger.warning(f"{symbol}价格无效: {last_price}")
            return False
        
        # 计算开仓手数
        lots = calculate_lots(api, symbol, direction)
        
        # 下单
        order = api.insert_order(
            symbol=symbol,
            direction=direction,
            offset='OPEN',
            volume=lots,
            price=last_price,
            limit_price_type='LIMIT'
        )
        
        # 等待订单成交或超时
        timeout = 30  # 30秒超时
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            api.wait_update()
            
            if order['status'] == 'FINISHED':
                # 订单成交
                fill_price = order.get('fill_price', 0)
                fill_volume = order.get('fill_volume', 0)
                
                console.print(f"[交易] {symbol} {direction} {fill_volume}手 价格: {fill_price}", style="trade")
                logger.info(f"{symbol}开仓成功: {direction} {fill_volume}手 价格: {fill_price}")
                
                # 记录交易
                insert_trade_record({
                    'product_name': symbol,
                    'direction': direction,
                    'offset': 'OPEN',
                    'volume': fill_volume,
                    'price': fill_price,
                    'reason': '自动开仓'
                })
                
                return True
            elif order['status'] == 'REJECTED':
                # 订单被拒绝
                logger.error(f"{symbol}订单被拒绝: {order.get('error_msg', '未知原因')}")
                return False
        
        # 超时，撤单
        api.cancel_order(order)
        logger.warning(f"{symbol}订单超时未成交")
        return False
    except Exception as e:
        logger.error(f"执行{symbol}开仓操作时出错: {str(e)}")
        return False


def monitor_market(api):
    """监控市场并执行交易策略"""
    console.print("开始监控市场...", style="info")
    
    while True:
        try:
            # 检查是否开启自动交易
            if get_config_value('open_auto', 'on') != 'on':
                console.print("自动交易已关闭，等待开启...", style="warning")
                time.sleep(30)
                continue
            
            # 遍历监控的品种
            for symbol in MONITOR_SYMBOLS:
                try:
                    # 这里应该实现更复杂的交易策略
                    # 以下是简单的随机交易示例
                    if random.random() < 0.05:  # 5%的概率触发交易
                        direction = 'BUY' if random.random() > 0.5 else 'SELL'
                        console.print(f"检测到{symbol}交易信号，方向: {direction}", style="info")
                        
                        # 执行开仓操作
                        open_position_non_shfe(api, symbol, direction)
                except Exception as e:
                    logger.error(f"监控{symbol}时出错: {str(e)}")
            
            # 获取监控间隔
            monitor_interval = int(get_config_value('monitor_interval', '30'))
            time.sleep(monitor_interval)
        except Exception as e:
            logger.error(f"监控市场时出错: {str(e)}")
            time.sleep(10)


def main():
    """主函数"""
    console.print("===== 自动交易程序启动 =====", style="info")
    console.print(f"启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", style="info")
    
    # 确保数据库表存在
    ensure_db_tables_exists()
    
    # 连接到TQSDK
    api = connect_to_tqsdk()
    
    if api:
        try:
            # 开始监控市场
            monitor_market(api)
        except KeyboardInterrupt:
            console.print("用户中断程序", style="info")
        except Exception as e:
            console.print(f"程序运行出错: {str(e)}", style="error")
            logger.error(f"程序运行出错: {str(e)}")
        finally:
            # 关闭TQSDK连接
            api.close()
            console.print("TQSDK连接已关闭", style="info")
    else:
        # 模拟模式运行
        console.print("进入模拟交易模式", style="warning")
        
        # 创建模拟API对象
        class MockApi:
            def get_position(self, symbol):
                return {'volume_long': 0, 'volume_short': 0}
            
            def get_order(self, symbol):
                return {}
            
            def get_quote(self, symbol):
                return {'last_price': 1000 + random.randint(-100, 100)}
            
            def get_instrument(self, symbol):
                return {'price_tick': 1, 'volume_multiple': 10}
            
            def insert_order(self, **kwargs):
                # 模拟80%的成功率
                if random.random() < 0.8:
                    return {
                        'status': 'FINISHED',
                        'fill_price': kwargs['price'],
                        'fill_volume': kwargs['volume']
                    }
                else:
                    return {'status': 'REJECTED', 'error_msg': '模拟订单拒绝'}
            
            def wait_update(self):
                time.sleep(0.1)
            
            def cancel_order(self, order):
                pass
            
            def close(self):
                pass
        
        mock_api = MockApi()
        
        try:
            # 开始模拟监控和交易
            console.print("开始模拟交易监控...", style="info")
            
            while True:
                # 检查是否开启自动交易
                if get_config_value('open_auto', 'on') != 'on':
                    console.print("自动交易已关闭，等待开启...", style="warning")
                    time.sleep(30)
                    continue
                
                # 遍历监控的品种
                for symbol in MONITOR_SYMBOLS:
                    try:
                        # 模拟交易信号
                        if random.random() < 0.2:  # 20%的概率触发交易（模拟模式提高概率）
                            direction = 'BUY' if random.random() > 0.5 else 'SELL'
                            console.print(f"[模拟] 检测到{symbol}交易信号，方向: {direction}", style="info")
                            
                            # 模拟开仓操作
                            # 获取当前行情
                            quote = mock_api.get_quote(symbol)
                            last_price = quote.get('last_price', 0)
                            
                            # 计算开仓手数
                            lots = calculate_lots(mock_api, symbol, direction)
                            
                            # 下单
                            order = mock_api.insert_order(
                                symbol=symbol,
                                direction=direction,
                                offset='OPEN',
                                volume=lots,
                                price=last_price,
                                limit_price_type='LIMIT'
                            )
                            
                            if order['status'] == 'FINISHED':
                                # 订单成交
                                fill_price = order.get('fill_price', 0)
                                fill_volume = order.get('fill_volume', 0)
                                
                                console.print(f"[模拟交易] {symbol} {direction} {fill_volume}手 价格: {fill_price}", style="trade")
                                logger.info(f"[模拟] {symbol}开仓成功: {direction} {fill_volume}手 价格: {fill_price}")
                                
                                # 记录交易
                                insert_trade_record({
                                    'product_name': symbol,
                                    'direction': direction,
                                    'offset': 'OPEN',
                                    'volume': fill_volume,
                                    'price': fill_price,
                                    'reason': '模拟开仓'
                                })
                    except Exception as e:
                        logger.error(f"模拟监控{symbol}时出错: {str(e)}")
                
                # 获取监控间隔
                monitor_interval = int(get_config_value('monitor_interval', '30'))
                time.sleep(monitor_interval)
        except KeyboardInterrupt:
            console.print("用户中断程序", style="info")
        except Exception as e:
            console.print(f"模拟模式运行出错: {str(e)}", style="error")
            logger.error(f"模拟模式运行出错: {str(e)}")
    
    console.print("程序已退出", style="info")


if __name__ == "__main__":
    main()